#include "lib/print.h"
#include "trap/trap.h"
#include "proc/cpu.h"
#include "mem/vmem.h"
#include "memlayout.h"
#include "common.h"
#include "riscv.h"
#include "syscall/syscall.h"
#include "mem/pmem.h"
#include "lib/str.h"

// in trampoline.S
extern char trampoline[];      // 内核和用户切换的代码
extern char user_vector[];     // 用户触发trap进入内核
extern char user_return[];     // trap处理完毕返回用户

// in trap.S
extern char kernel_vector[];   // 内核态trap处理流程

// in trap_kernel.c
extern char* interrupt_info[16]; // 中断错误信息
extern char* exception_info[16]; // 异常错误信息

static void user_page_fault_handler(proc_t *p, uint64 scause, uint64 stval)
{
    // stval 存储了导致异常的虚拟地址
    uint64 va = stval;

    void* pa = pmem_alloc(false);
    if(pa == 0){
        // 如果内核无法分配物理页，只能杀掉进程
        printf("user_page_fault_handler: kalloc() failed\n");
    }

    va = ALIGN_DOWN(va,PGSIZE);

    memset(pa, 0, PGSIZE);

    vm_mappages(p->pgtbl, va, PGSIZE, (uint64)pa, PTE_U | PTE_R | PTE_W);

}


// 在user_vector()里面调用
// 用户态trap处理的核心逻辑
void trap_user_handler()
{
    uint64 sepc = r_sepc();          // 记录了发生异常时的pc值
    uint64 sstatus = r_sstatus();    // 与特权模式和中断相关的状态信息
    uint64 scause = r_scause();      // 引发trap的原因
    uint64 stval = r_stval();        // 发生trap时保存的附加信息(不同trap不一样)
    proc_t* p = myproc();
    // 确认trap来自U-mode
    assert((sstatus & SSTATUS_SPP) == 0, "trap_user_handler: not from u-mode");

    int trap_id = scause & 0xf;  // 获取中断或异常的类型

    // 将中断和异常发送到 kernel_trap()
    w_stvec((uint64)kernel_vector);

    // 设置 S Exception Program Counter 为保存的用户 PC
    p->tf->epc = sepc;

    switch (trap_id) 
    {
        case 8: // 环境调用 (系统调用)
            p->tf->epc += 4;
            intr_on();
            syscall();
            break;
        case 1: // S-mode software interrupt
            timer_interrupt_handler();
            break;
        case 5: // S-mode timer interrupt
            timer_interrupt_handler();
            break;
        case 9: // S-mode external interrupt
            external_interrupt_handler();
            break;
        case 13: // Store Page Fault
        {
            uint64 va = stval;
            if(cow_alloc(p->pgtbl, va) != 0)
                printf("cow_alloc(p->pagetable, va) != 0");
            break;
        }
        case 15: // Load/Instruction Fetch Page Fault
        {
        user_page_fault_handler(p, scause, stval);
        break;
        }
        
        // 添加其他中断或异常的处理
        default:
            if (scause & (1L << 63)) 
            {
                // 中断处理
                if (trap_id < 16) {
                    printf("Interrupt: %s\n", interrupt_info[trap_id]);
                    printf("trap occurred at sepc=%p, stval=%p\n", sepc, stval);
                    panic("trap_user_handler");
                } else {
                    printf("Unexpected interrupt: trap_id=%d\n", trap_id);
                    printf("trap occurred at sepc=%p, stval=%p\n", sepc, stval);
                    panic("trap_user_handler");
                }
            } 
            else {
                // 异常处理
                if (trap_id < 16) {
                    printf("Exception: %s\n", exception_info[trap_id]);
                    printf("trap occurred at sepc=%p, stval=%p\n", sepc, stval);
                    panic("trap_user_handler");
                } else {
                    printf("Unexpected exception: trap_id=%d\n", trap_id);
                    printf("trap occurred at sepc=%p, stval=%p\n", sepc, stval);
                    panic("trap_user_handler");
                }
            }
            break;
    }

    // 完成 trap 处理，准备返回用户态
    trap_user_return();
}

// 调用user_return()
// 内核态返回用户态
void trap_user_return()
{
    proc_t *p = myproc();          // 获取当前用户进程
    // 我们即将将陷阱处理的目标从 kerneltrap() 切换到 usertrap()，
    // 所以在回到用户空间之前关闭中断。
    intr_off();

    w_stvec(TRAMPOLINE + (user_vector - trampoline));

    // 设置 uservec 需要的陷阱帧值，并恢复用户寄存器
    p->tf->kernel_satp = r_satp();         // 内核页表
    p->tf->kernel_sp = p->kstack + PGSIZE; // 进程的内核栈
    p->tf->kernel_trap = (uint64)trap_user_handler;
    p->tf->kernel_hartid = r_tp(); // hartid 用于 CPU ID

    unsigned long x = r_sstatus();
    x &= ~SSTATUS_SPP; 
    x |= SSTATUS_SPIE;
    w_sstatus(x);

    w_sepc(p->tf->epc);

    uint64 satp = MAKE_SATP(p->pgtbl);

    uint64 fn = TRAMPOLINE + (user_return - trampoline);
    ((void (*)(uint64, uint64))fn)(TRAPFRAME, satp);
}